from abc import abstractmethod
from typing import TypeVar, Generic, List, Union, Dict, Optional, Type

from bson import ObjectId
from motor.core import AgnosticCollection
from pydantic import BaseModel, ValidationError
from pymongo import InsertOne, UpdateOne, UpdateMany, ReplaceOne, DeleteOne, DeleteMany

from common.singleton import Singleton
from config import logger

T = TypeVar('T', bound=BaseModel)
OperationType = Union[InsertOne, UpdateOne, UpdateMany, ReplaceOne, DeleteOne, DeleteMany]


class BaseDAO(Generic[T], Singleton):

    @abstractmethod
    def get_db(self) -> AgnosticCollection:
        pass

    async def insert_one(self, data: T):
        db = self.get_db()
        insert_data = data.dict(by_alias=True)
        result = await db.insert_one(insert_data)
        data.id = result.inserted_id
        return result.inserted_id

    async def insert_many(self, data: List[T]):
        db = self.get_db()
        inserts = []
        for d in data:
            insert_data = d.dict(by_alias=True)
            inserts.append(insert_data)
        result = await db.insert_many(inserts)
        for _id, d in zip(result.inserted_ids, data):
            d.id = _id
        return result.inserted_ids

    async def insert_or_update_by_id(self, _id: Union[str, ObjectId], data: T):
        db = self.get_db()
        if isinstance(_id, str):
            _id = ObjectId(_id)
        return await db.replace_one({'_id': _id}, data.dict(by_alias=True), upsert=True)

    async def insert_or_update_by_query(self, query: Dict, data: T):
        db = self.get_db()
        return await db.replace_one(query, data.dict(by_alias=True), upsert=True)

    async def update_by_id(self, id_: Union[str, ObjectId], update_data: Dict):
        db = self.get_db()
        if isinstance(id_, str):
            id_ = ObjectId(id_)
        query = {
            '_id': id_
        }

        update_ = {
            '$set': update_data
        }
        return await db.update_one(query, update_)

    async def update_many_by_query(self, query: Dict, update_data: Dict):
        db = self.get_db()
        update_ = {
            '$set': update_data
        }
        return await db.update_many(query, update_)

    async def get_by_id(self, id_: Union[str, ObjectId]) -> Optional[T]:
        db = self.get_db()
        if isinstance(id_, str):
            id_ = ObjectId(id_)
        query = {
            '_id': id_
        }
        res: dict = await db.find_one(query)
        if res:
            obj = self.dict_to_model(res)
            return obj
        return

    async def get_one_by_query(self, query: Dict, **kwargs) -> Optional[T]:
        db = self.get_db()

        res: dict = await db.find_one(query, **kwargs)
        if res:
            obj = self.dict_to_model(res)
            return obj
        return

    async def get_by_query(self, query: Dict, **kwargs) -> List[T]:
        db = self.get_db()
        res = []
        async for item in db.find(query, **kwargs):
            obj = self.dict_to_model(item)
            res.append(obj)
        return res

    async def get_by_query_params(self, query_params: Dict) -> List[T]:
        db = self.get_db()
        query = {}
        for k, v in query_params.items():
            if v:
                query[k] = v
        res = [i async for i in db.find(query)]
        return res

    async def delete_by_id(self, id_: Union[str, ObjectId]):
        db = self.get_db()
        if isinstance(id_, str):
            id_ = ObjectId(id_)
        query = {
            '_id': id_
        }
        return await db.delete_one(query)

    async def delete_many_by_query(self, query: Dict):
        db = self.get_db()
        return await db.delete_many(query)

    async def bulk_write(self, data: List[OperationType]):
        db = self.get_db()
        return await db.bulk_write(data)

    def _get_generate_class(self) -> Type[T]:
        # __orig_bases__ 是一个特殊的属性，通常只有在使用 Generic 类型时才会出现，特别是当你用 typing 模块定义泛型类时。它存储了类的原始基类（包括泛型类型参数）。
        # 对于普通的类，__orig_bases__ 是不存在的
        base_cls = self.__orig_bases__[0]
        cls = base_cls.__args__[0]
        return cls

    def dict_to_model(self, data: Dict) -> T:
        cls = self._get_generate_class()
        try:
            cls_obj = cls(**data)
            return cls_obj
        except ValidationError as e:
            logger.warning(f'mongo result to model error: {e.errors()}')
            raise
